# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""RMSprop optimizer implementation."""

from keras.optimizers.optimizer_experimental import optimizer
from keras.utils import generic_utils
import tensorflow.compat.v2 as tf


@generic_utils.register_keras_serializable()
class RMSprop(optimizer.Optimizer):
  r"""Optimizer that implements the RMSprop algorithm.

  The gist of RMSprop is to:

  - Maintain a moving (discounted) average of the square of gradients
  - Divide the gradient by the root of this average

  This implementation of RMSprop uses plain momentum, not Nesterov momentum.

  The centered version additionally maintains a moving average of the
  gradients, and uses that average to estimate the variance.

  Attributes:
    learning_rate: Initial value for the learning rate:
      either a floating point value,
      or a `tf.keras.optimizers.schedules.LearningRateSchedule` instance.
      Defaults to 0.001.
    rho: float, defaults to 0.9. Discounting factor for the old gradients.
    momentum: float, defaults to 0.0. If not 0.0., the optimizer tracks the
      momentum value, with a decay rate equals to `1 - momentum`.
    epsilon: A small constant for numerical stability. This epsilon is
      "epsilon hat" in the Kingma and Ba paper (in the formula just before
      Section 2.1), not the epsilon in Algorithm 1 of the paper. Defaults to
      1e-7.
    centered: Boolean. If `True`, gradients are normalized by the estimated
      variance of the gradient; if False, by the uncentered second moment.
      Setting this to `True` may help with training, but is slightly more
      expensive in terms of computation and memory. Defaults to `False`.
    clipnorm: see the `clipnorm` argument of `optimizer_experimental.Optimizer`.
    clipvalue: see the `clipvalue` argument of
      `optimizer_experimental.Optimizer`.
    global_clipnorm: see the `global_clipnorm` argument of
      `optimizer_experimental.Optimizer`.
    use_ema: see the `use_ema` argument of `optimizer_experimental.Optimizer`.
    ema_momentum: see the `ema_momentum` argument of
      `optimizer_experimental.Optimizer`.
    ema_overwrite_frequency: see the `ema_overwrite_frequency` argument of
      `optimizer_experimental.Optimizer`.
    jit_compile: see the `jit_compile` argument of
      `optimizer_experimental.Optimizer`.
    name: Optional name prefix for the operations created when applying
      gradients. Defaults to `"RMSprop"`.
    **kwargs: see the `**kwargs` argument of `optimizer_experimental.Optimizer`.

  Usage:

  >>> opt = tf.keras.optimizers.RMSprop(learning_rate=0.1)
  >>> var1 = tf.Variable(10.0)
  >>> loss = lambda: (var1 ** 2) / 2.0    # d(loss) / d(var1) = var1
  >>> step_count = opt.minimize(loss, [var1]).numpy()
  >>> var1.numpy()
  9.683772

  Reference:
    - [Hinton, 2012](
      http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)

  """

  def __init__(self,
               learning_rate=0.001,
               rho=0.9,
               momentum=0.0,
               epsilon=1e-7,
               centered=False,
               clipnorm=None,
               clipvalue=None,
               global_clipnorm=None,
               use_ema=False,
               ema_momentum=0.99,
               ema_overwrite_frequency=100,
               jit_compile=False,
               name='RMSprop',
               **kwargs):
    super(RMSprop, self).__init__(
        clipnorm=clipnorm,
        clipvalue=clipvalue,
        global_clipnorm=global_clipnorm,
        use_ema=use_ema,
        ema_momentum=ema_momentum,
        ema_overwrite_frequency=ema_overwrite_frequency,
        jit_compile=jit_compile,
        name=name,
        **kwargs)
    self._learning_rate = self._build_learning_rate(learning_rate)
    self.rho = rho
    self.momentum = momentum
    self.epsilon = epsilon
    self.centered = centered

  def build(self, var_list):
    super().build(var_list)
    if hasattr(self, '_built') and self._built:
      return
    self._built = True

    self._velocities = []
    for var in var_list:
      self._velocities.append(
          self.add_variable_from_reference(var, 'velocity'))

    self._momentums = []
    if self.momentum > 0:
      for var in var_list:
        self._momentums.append(
            self.add_variable_from_reference(var, 'momentum'))

    self._average_gradients = []
    if self.centered:
      for var in var_list:
        self._average_gradients.append(
            self.add_variable_from_reference(var, 'average_gradient'))

  def update_step(self, gradient, variable):
    """Update step given gradient and the associated model variable."""
    if self._var_key(variable) not in self._index_dict:
      raise KeyError(f'Optimizer cannot recognize variable {variable.name}, '
                     f'this usually means you are calling an optimizer '
                     f'previously used on a different model. Please try '
                     f'creating a new optimizer instance.')
    lr = tf.cast(self.learning_rate, variable.dtype)

    var_key = self._var_key(variable)
    velocity = self._velocities[self._index_dict[var_key]]
    momentum = None
    if self.momentum > 0:
      momentum = self._momentums[self._index_dict[var_key]]
    average_grad = None
    if self.centered:
      average_grad = self._average_gradients[self._index_dict[var_key]]

    rho = self.rho

    if isinstance(gradient, tf.IndexedSlices):
      # Sparse gradients.
      velocity.assign(rho * velocity)
      velocity.scatter_add(tf.IndexedSlices(
          tf.square(gradient.values) * (1 - rho), gradient.indices))
      if self.centered:
        average_grad.assign(rho * average_grad)
        average_grad.scatter_add(
            tf.IndexedSlices(
                tf.square(gradient.values) * (1 - rho), gradient.indices))
        velocity.assign_add(-tf.square(average_grad))
      velocity_value = tf.gather(velocity, gradient.indices)
      transformed_grad = tf.IndexedSlices(
          gradient.values / (tf.sqrt(velocity_value) + self.epsilon),
          gradient.indices)

      if self.momentum > 0:
        momentum.assign(self.momentum * momentum)
        momentum.scatter_add(transformed_grad)
        variable.assign_add(-lr * momentum)
      else:
        variable.scatter_add(
            tf.IndexedSlices(-lr * transformed_grad.values,
                             transformed_grad.indices))
    else:
      # Dense gradients.
      velocity.assign(rho * velocity + (1 - rho) * tf.square(gradient))
      if self.centered:
        average_grad.assign(rho * average_grad +
                            (1 - rho) * tf.square(gradient))
        velocity.assign_add(-tf.square(average_grad))
      transformed_grad = gradient / (tf.sqrt(velocity) + self.epsilon)
      if self.momentum > 0:
        momentum.assign(self.momentum * momentum + transformed_grad)
        variable.assign_add(-lr * momentum)
      else:
        variable.assign_add(-lr * transformed_grad)

  def get_config(self):
    config = super(RMSprop, self).get_config()

    config.update({
        'learning_rate': self._serialize_hyperparameter(self._learning_rate),
        'rho': self.rho,
        'momentum': self.momentum,
        'epsilon': self.epsilon,
        'centered': self.centered,
    })
    return config
